#!/usr/bin/env Rscript
# sudo apt install clustalw phylip
library(tidyverse)
library(immunarch)

# test data ----
data("immdata")

# Visualise the length distribution of CDR3
repExplore(immdata$data, "lens") %>% vis()
# Visualise the relative abundance of clonotypes
repClonality(immdata$data, "homeo") %>% vis()  

# Build the heatmap of public clonotypes shared between repertoires
repOverlap(immdata$data) %>% vis()
# Visualise the V-gene distribution for the first repertoire
geneUsage(immdata$data[[1]]) %>% vis()
# Visualise the Chao1 diversity of repertoires, grouped by the patient status
repDiversity(immdata$data) %>%
  vis(.by = "Status", .meta = immdata$meta)  

gu <- geneUsage(immdata$data, .norm = TRUE)
immunr_hclust(t(as.matrix(gu[, -1])), .dist = FALSE)

# example BCR data -------
data("bcrdata")

#calulate distance matrix
distBCR <- bcrdata$data |>
  top(500) |>
  seqDist()

#find clusters
bcrdata$data <- bcrdata$data %>%
  top(500) |>
  seqCluster(distBCR, .perc_similarity = 0.6)

example_tibble <- bcrdata$data
example_tibble[[1]] |> count(Cluster) |> dim()

#germline example
bcrdata$data %>%
  top(1) %>%
  repGermline(.threads = 1) %>%
  .$full_clones %>%
  .$Germline.sequence

# Aligning sequences within a clonal lineage
data("bcrdata")
bcr_data <- bcrdata$data

# take clusters that contain at least 1 sequence
align_dt <- bcr_data %>%
  seqCluster(seqDist(bcr_data), .perc_similarity = 0.6) %>%
  repGermline(.threads = 1) %>%
  repAlignLineage(.min_lineage_sequences = 6, .align_threads = 2, .nofail = TRUE)

bcr_data %>%
  seqCluster(seqDist(bcr_data), .perc_similarity = 0.6) %>%
  repGermline(.threads = 1) -> t

t$full_clones |>
  count(Cluster) |>
  filter(n >= 6)

# Alignment of sequences from the first cluster
image(align_dt$full_clones$Alignment[[1]], grid = TRUE)

bcr <- align_dt %>%
  repClonalFamily(.threads = 2, .nofail = TRUE)
#plot visualization of the first tree
vis(bcr[["full_clones"]][["TreeStats"]][[1]])

#take sequence where number of AA mutations is not 0
f <- bcr[["full_clones"]][["TreeStats"]][[1]]
#rename these leaves
f[f$DistanceAA != 0, ]['Type'] = 'mutationAA'
#new tree
vis(f)

#get all clone IDs from align_dt
clone_ids <- unnest(align_dt[["full_clones"]], "Sequences")[["Clone.ID"]]
#run repClonalFamily with assigning some of these clones to differently named and colored groups
bcr_with_groups <- align_dt %>%
  repClonalFamily(.vis_groups = list(
    Group1 = clone_ids[1],
    Group2 = clone_ids[3],
    Group3 = list(clone_ids[5], clone_ids[2]),
    Group4 = c(clone_ids[7], clone_ids[4])
  ), .threads = 2, .nofail = TRUE
  )
#display the first tree from repClonalFamily results
vis(bcr_with_groups[["full_clones"]][["TreeStats"]][[1]])

# SHM analysis
shm_data <- bcr %>% repSomaticHypermutation(.threads = 2, .nofail = TRUE)

image(shm_data$full_clones$Germline.Alignment.V[[3]], grid = TRUE)

image(shm_data$full_clones$Germline.Alignment.J[[3]], grid = TRUE)

# estimate mutation rate
shm_data$full_clones %>%
  mutate(Mutation.Rate = Mutations / (nchar(Sequence) - nchar(CDR3.nt))) %>%
  select(Clone.ID, Mutation.Rate)

# my data --------
data("bcrdata")

write_csv(bcrdata$data[[1]], 'Archive/FPP_TRPV3/data/BCR/ref/immunarch_bcr_example.csv')

farne_migmap <- repLoad("Archive/FPP_TRPV3/data/BCR/")

farne_migmap$data[['VH-1.migmap']] |>
  write_csv('Archive/FPP_TRPV3/data/BCR/parsed/parsed_farn_VH.csv')

farne_migmap$data[['VK-1.migmap']] |>
  write_csv('Archive/FPP_TRPV3/data/BCR/parsed/parsed_farn_VK.csv')

# edit in excel
farne_bcr <- repLoad("Archive/FPP_TRPV3/data/BCR/parsed/")

vdj_tibble <- farne_bcr$data

# number of clones
vdj_tibble |>
  repExplore('volume') |>
  vis()

# number of clonotypes
vdj_tibble |>
  repExplore('clone') |>
  vis()

# length of CDR3
vdj_tibble |>
  repExplore('len', 'aa') |>
  vis() +
  scale_x_continuous(breaks = c(8, 10, 12, 14, 16, 18, 20))

# Visualise the relative abundance of clonotypes
repClonality(vdj_tibble, "clonal.prop") %>% vis()

# vdj gene usage analysis
vdj_tibble |>
  geneUsage('hs.ighv') |>
  vis() +
  ggtitle('IGV gene usage')

vdj_tibble |>
  geneUsage('hs.ighd') |>
  vis() +
  ggtitle('IGD gene usage')

vdj_tibble |>
  geneUsage('hs.ighj') |>
  vis() +
  ggtitle('IGJ gene usage')

# results of hs.ighv == hs.igkv, hs.ighj == hs.igvj?
vdj_tibble |>
  geneUsage('hs.igkj') |>
  vis()

vdj_tibble |>
  geneUsage('hs.ighv') |>
  immunr_hclust()

vdj_tibble[[1]] |>
  spectratype('count', 'aa+v') |>
  ggplot(aes(Length, Val, fill = Gene)) +
  geom_col()

# Kmer analysis ------
vdj_tibble[[1]] |>
  getKmers(.k = 13, .col = 'aa') |>
  slice_max(Count, n = 12) |>
  mutate(`5-kmer AA motif` = fct_reorder(Kmer, Count)) |>
  ggplot(aes(`5-kmer AA motif`, Count)) +
  geom_col() +
  coord_flip() +
  theme_classic()

vdj_tibble[[1]] |>
  getKmers(.k = 13, .col = 'aa') |>
  kmer_profile() |>
  vis(.plot = 'seq')

vdj_tibble[[2]] |>
  getKmers(.k = 5, .col = 'aa') |>
  slice_max(Count, n = 12) |>
  mutate(`5-kmer AA motif` = fct_reorder(Kmer, Count)) |>
  ggplot(aes(`5-kmer AA motif`, Count)) +
  geom_col() +
  coord_flip() +
  theme_classic()

vdj_tibble[[2]] |>
  getKmers(.k = 11, .col = 'aa') |>
  kmer_profile() |>
  vis(.plot = 'seq')

# BCR analysis of my data ---------
#calulate distance matrix
distBCR <- vdj_tibble |>
  seqDist()

#find clusters
farne_clustered <- vdj_tibble |>
  seqCluster(distBCR, .perc_similarity = 0.6)

farne_clustered$parsed_farn_VH |>
  ggplot(aes(Cluster)) + geom_bar() + coord_flip()

farne_clustered$parsed_farn_VK |>
  ggplot(aes(Cluster)) + geom_bar() + coord_flip()

# Aligning sequences within a clonal lineage
aligned_gl <- farne_clustered$parsed_farn_VK %>%
  repGermline(.threads = 1) %>%
  repAlignLineage(.align_threads = 2, .min_lineage_sequences = 2)

# Alignment of sequences from the first cluster
image(aligned_gl$immunarch_parsed_farn$Alignment[[2]], grid = TRUE)

farne_family <- aligned_gl %>%
  repClonalFamily(.threads = 3)

# plot visualization of the first tree
vis(bcr[["full_clones"]][["TreeStats"]][[4]])
bcr$full_clones$TreeStats
#take sequence where number of AA mutations is not 0
f <- bcr[["full_clones"]][["TreeStats"]][[1]]
#rename these leaves
f[f$DistanceAA != 0, ]['Type'] = 'mutationAA'
#new tree
vis(f)

#get all clone IDs from align_dt
clone_ids <- unnest(align_dt[["full_clones"]], "Sequences")[["Clone.ID"]]
#run repClonalFamily with assigning some of these clones to differently named and colored groups
bcr_with_groups <- align_dt %>%
  repClonalFamily(.vis_groups = list(
    Group1 = clone_ids[1],
    Group2 = clone_ids[3],
    Group3 = list(clone_ids[5], clone_ids[2]),
    Group4 = c(clone_ids[7], clone_ids[4])
  ), .threads = 2, .nofail = TRUE
  )
#display the first tree from repClonalFamily results
vis(bcr_with_groups[["full_clones"]][["TreeStats"]][[2]])

# SHM analysis
bcr_data <- bcrdata$data

shm_data <- bcr %>% repSomaticHypermutation(.threads = 2, .nofail = TRUE)

image(shm_data$full_clones$Germline.Alignment.V[[3]], grid = TRUE)

image(shm_data$full_clones$Germline.Alignment.J[[3]], grid = TRUE)

# estimate mutation rate
shm_data$full_clones %>%
  mutate(Mutation.Rate = Mutations / (nchar(Sequence) - nchar(CDR3.nt))) %>%
  select(Clone.ID, Mutation.Rate)

